// langchain loader 是RAG的基础功能 txt，pdf，excel....
// 加载网页内容
import {
  PuppeteerWebBaseLoader
} from '@langchain/community/document_loaders/web/puppeteer';
import {
  RecursiveCharacterTextSplitter
} from 'langchain/text_splitter';
import { createOpenAI } from '@ai-sdk/openai';
import {
  embed // 向量嵌入
} from "ai";
import {
  config
} from 'dotenv';
config(); // 加载环境变量
import { createClient } from '@supabase/supabase-js';

const supabase = createClient(
  process.env.SUPABASE_URL ?? "",
  process.env.SUPABASE_KEY ?? ""
);

const openai = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,     // 自定义密钥
  baseURL: process.env.OPENAI_API_BASE_URL, // 如使用代理或 Azure
});

// 优化后的分块配置
const splitter = new RecursiveCharacterTextSplitter({
  chunkSize: 1000, // 适当增大块大小以包含更多上下文
  chunkOverlap: 200, // 增加重叠区域以保持语义连贯
  separators: ['\n\n', '\n', '。', '！', '？', '；', '，', ' ', ''], // 中文友好的分隔符
});

const scrapePage = async (url: string): Promise<string> => {
  const loader = new PuppeteerWebBaseLoader(url, {
    launchOptions: {
      executablePath: 'C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe',
      headless: true,
    },
    gotoOptions: {
      waitUntil: 'domcontentloaded',
    },
    evaluate: async (page, browser) => {
      const result = await page.evaluate(() => document.body.innerText || document.body.textContent || '');
      await browser.close();
      return result;
    }
  });
  
  const scrapedContent = await loader.scrape();
  // 更彻底的清理，保留更多有用的文本
  return scrapedContent
    .replace(/<script\b[^<]*(?:(?!<\/script>)<[^<]*)*<\/script>/gi, '') // 移除脚本
    .replace(/<style\b[^<]*(?:(?!<\/style>)<[^<]*)*<\/style>/gi, '') // 移除样式
    .replace(/<[^>]*>?/gm, ' ') // 用空格替换标签，避免单词粘连
    .replace(/\s+/g, ' ') // 合并多个空白字符
    .trim();
}

// 优化的分块处理函数
const processContentWithBetterChunking = async (content: string, url: string) => {
  // 1. 首先按段落分割
  const paragraphs = content.split(/\n\s*\n/).filter(p => p.trim().length > 0);
  
  const chunks = [];
  
  for (const paragraph of paragraphs) {
    const paragraphText = paragraph.trim();
    
    if (paragraphText.length <= 800) {
      // 如果段落较短，直接作为一个块
      if (paragraphText.length > 50) { // 过滤掉太短的内容
        chunks.push(paragraphText);
      }
    } else {
      // 如果段落较长，使用递归分块器
      const paragraphChunks = await splitter.splitText(paragraphText);
      chunks.push(...paragraphChunks);
    }
  }
  
  // 2. 合并过小的块
  const mergedChunks: string[] = [];
  let currentChunk = '';
  
  for (const chunk of chunks) {
    if (currentChunk.length + chunk.length <= 1000) {
      currentChunk += (currentChunk ? '\n\n' : '') + chunk;
    } else {
      if (currentChunk) {
        mergedChunks.push(currentChunk);
      }
      currentChunk = chunk;
    }
  }
  
  if (currentChunk) {
    mergedChunks.push(currentChunk);
  }
  
  // 3. 处理最终块
  const finalChunks: string[] = [];
  for (const chunk of mergedChunks) {
    if (chunk.length > 100) { // 过滤掉太短的块
      finalChunks.push(chunk);
    }
  }
  
  console.log(`URL: ${url} - 原始内容长度: ${content.length}, 生成块数: ${finalChunks.length}`);
  
  return finalChunks;
}

const loadData = async (webpages: string[]) => {
  for (const url of webpages) {
    try {
      console.log(`开始处理: ${url}`);
      
      const content = await scrapePage(url);
      
      if (!content || content.length < 100) {
        console.warn(`页面内容过少或为空: ${url}`);
        continue;
      }
      
      const chunks = await processContentWithBetterChunking(content, url);
      
      console.log(`为 ${url} 生成了 ${chunks.length} 个块`);
      
      let successCount = 0;
      let errorCount = 0;
      
      // 批量处理嵌入和存储
      for (let i = 0; i < chunks.length; i++) {
        const chunk = chunks[i];
        
        try {
          const { embedding } = await embed({
            model: openai.embedding('text-embedding-3-small'),
            value: chunk
          });
          
          // 只插入数据库表中实际存在的字段
          const { error } = await supabase.from('chunks').insert({
            content: chunk,
            vector: embedding,
            url: url,
          });
          
          if (error) {
            console.error(`插入块 ${i} 时出错:`, error);
            errorCount++;
          } else {
            successCount++;
            if (successCount % 10 === 0) {
              console.log(`已成功插入 ${successCount} 个块`);
            }
          }
        } catch (embedError) {
          console.error(`为块 ${i} 生成嵌入时出错:`, embedError);
          errorCount++;
        }
      }
      
      console.log(`完成 ${url}: 成功 ${successCount}, 失败 ${errorCount}`);
      
    } catch (error) {
      console.error(`处理 URL ${url} 时发生错误:`, error);
    }
  }
  
  console.log('所有页面处理完成');
}

// 知识库的来源，可配置
loadData([
  "https://prts.wiki/w/%E4%B8%B0%E5%B7%9D%E7%A5%A5%E5%AD%90"
]);